// SPDX-License-Identifier: GPL-2.0-only
/*
 * kretprobe_example.c
 *
 * Here's a sample kernel module showing the use of return probes to
 * report the return value and total time taken for probed function
 * to run.
 *
 * usage: insmod kretprobe_example.ko func=<func_name>
 *
 * If no func_name is specified, _do_fork is instrumented
 *
 * For more information on theory of operation of kretprobes, see
 * Documentation/kprobes.txt
 *
 * Build and insert the kernel module as done in the kprobe example.
 * You will see the trace data in /var/log/messages and on the console
 * whenever the probed function returns. (Some messages may be suppressed
 * if syslogd is configured to eliminate duplicate messages.)
 */

#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/kprobes.h>
#include <linux/ktime.h>
#include <linux/limits.h>
#include <linux/sched.h>
#include "kprobe_prehook.h"
#include "data_aware.h"
#include "hook.h"

#ifdef TK5
static char func_name[NAME_MAX] = "x64_sys_call";
#else
static char func_name[NAME_MAX] = "do_syscall_64";
#endif
char symbol_kret_new[NAME_MAX];
char func_name_new[NAME_MAX];

#ifdef CONFIG_X86_64
/******stat_do_syscall_64: only for test********/
void stat_do_syscall_64(unsigned long nr, struct pt_regs *regs)
{
	do_syscall_64(nr, regs);
}
#endif

/* Here we use the entry_hanlder to timestamp function entry */
static int entry_handler(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct kret_data *data;

	if (!current->mm)
		return 1;	/* Skip kernel threads */

	current->numa_faults_locality[0] = 0;
	current->numa_faults_locality[1] = 0;
	current->numa_faults_locality[2] = 0;

	data = (struct kret_data *)ri->data;

#ifdef CONFIG_X86_64
	set_sys_nr(regs->ax);
	data->nr = regs->ax;
#elif defined CONFIG_ARM64
	set_sys_nr(regs->syscallno);
	data->nr = regs->syscallno;
#endif
	stat_stat_syscall_enter(data);
#if 0
	pr_info("kret sys:%ld", regs->ax);
	data = (struct kret_data *)ri->data;
	data->entry_stamp = ktime_get();
#endif
	return 0;
}

static int new_entry_handler(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct kret_data *data = (struct kret_data *)ri->data;

	stat_func_enter(data);
	return 0;
}

/*
 * Return-probe handler: Log the return value and duration. Duration may turn
 * out to be zero consistently, depending upon the granularity of time
 * accounting on the platform.
 */
static int ret_handler(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct kret_data *data = (struct kret_data *)ri->data;

	stat_stat_syscall_exit(get_sys_nr(), data);
	current->numa_faults_locality[0] = 0;
	current->numa_faults_locality[1] = 0;
	current->numa_faults_locality[2] = 0;
#if 0
	unsigned long retval = regs_return_value(regs);
	struct kret_data *data = (struct kret_data *)ri->data;
	s64 delta;
	ktime_t now;

	now = ktime_get();
	delta = ktime_to_ns(ktime_sub(now, data->entry_stamp));
	pr_info("%s returned %lu and took %lld ns to execute:%ld\n",
			func_name, retval, (long long)delta, regs->ax);
#endif
	return 0;
}

static int new_ret_handler(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct kret_data *data = (struct kret_data *)ri->data;

	stat_func_exit(data);
	return 0;
}
static struct kretprobe resident_kretprobe = {
	.handler		= ret_handler,
	.entry_handler		= entry_handler,
	.data_size		= sizeof(struct kret_data),
	/* Probe up to 20 instances concurrently. */
	.maxactive		= 20,
};
static struct kretprobe new_kretprobe = {
	.handler		= new_ret_handler,
	.entry_handler		= new_entry_handler,
	.data_size		= sizeof(struct kret_data),
	/* Probe up to 20 instances concurrently. */
	.maxactive		= 200,
};
/* for one function hook */
static int register_kret_new_func(char *name)
{
	int ret;

	strcpy(func_name_new, name);
	new_kretprobe.kp.symbol_name = func_name_new;
	new_kretprobe.handler = new_ret_handler;
	new_kretprobe.entry_handler = new_entry_handler;
	new_kretprobe.data_size = sizeof(struct kret_data);
	new_kretprobe.maxactive = 200;
	ret = register_kretprobe(&new_kretprobe);
	if (ret < 0) {
		pr_err("register_kretprobe failed, returned %d, %s, %s\n", ret, name, func_name_new);
		return ret;
	}
	pr_info("Planted return probe at %s: %p\n",
			new_kretprobe.kp.symbol_name, new_kretprobe.kp.addr);
	return 0;
}

static void ungister__kret_new_func(void)
{
	pr_info("unregister_kretprobe %s\n", func_name_new);
	unregister_kretprobe(&new_kretprobe);
	memset(&new_kretprobe, 0, sizeof(new_kretprobe));
}
/* for hot path batch process */
#define KRET_FUNC(index) \
char func_name_new_##index[NAME_MAX]; \
static int new_entry_handler_##index(struct kretprobe_instance *ri, struct pt_regs *regs) \
{ \
	struct kret_data *data = (struct kret_data *)ri->data;\
\
	save_start_time(index);\
	return 0;\
}\
static int new_ret_handler_##index(struct kretprobe_instance *ri, struct pt_regs *regs)\
{\
	save_total_time(index);\
	return 0;\
}\
static struct kretprobe new_kretprobe_##index = {\
	.handler		= new_ret_handler_##index,\
	.entry_handler		= new_entry_handler_##index,\
	.data_size		= sizeof(struct kret_data),\
	.maxactive		= 200,\
};\
static int register_kret_new_func_##index(char *func_name)\
{\
	int ret;\
\
	strcpy(func_name_new_##index, func_name);\
	new_kretprobe_##index.kp.symbol_name = func_name_new_##index;\
	new_kretprobe_##index.handler = new_ret_handler_##index;\
	new_kretprobe_##index.entry_handler = new_entry_handler_##index;\
	new_kretprobe_##index.data_size = sizeof(struct kret_data);\
	new_kretprobe_##index.maxactive = 200;\
	ret = register_kretprobe(&new_kretprobe_##index);\
	if (ret < 0) {\
		pr_err("register_kretprobe failed, returned %d, %s, %s\n", ret, func_name, func_name_new_##index);\
		return ret;\
	}\
	pr_info("Planted return probe at %s: %p\n", new_kretprobe_##index.kp.symbol_name, new_kretprobe_##index.kp.addr); \
	return 0;\
}\
static void unregister_kret_new_func_##index(void)\
{\
	pr_info("unregister_kretprobe %s\n", func_name_new_##index);\
	unregister_kretprobe(&new_kretprobe_##index);\
	memset(&new_kretprobe, 0, sizeof(new_kretprobe));\
}
KRET_FUNC(1)
KRET_FUNC(2)
KRET_FUNC(3)
KRET_FUNC(4)
KRET_FUNC(5)
KRET_FUNC(6)
KRET_FUNC(7)
KRET_FUNC(8)
KRET_FUNC(9)
KRET_FUNC(10)
KRET_FUNC(11)
KRET_FUNC(12)
KRET_FUNC(13)
KRET_FUNC(14)
KRET_FUNC(15)
KRET_FUNC(16)
KRET_FUNC(17)
KRET_FUNC(18)
KRET_FUNC(19)
KRET_FUNC(20)
KRET_FUNC(21)
KRET_FUNC(22)
KRET_FUNC(23)
KRET_FUNC(24)
KRET_FUNC(25)
KRET_FUNC(26)
KRET_FUNC(27)
KRET_FUNC(28)
KRET_FUNC(29)
KRET_FUNC(30)

int register_kret_new_func_batch(int index, char* func_name)
{
	int ret = -1;

	pr_info("----%s, %d", func_name, index);
	switch(index) {
	case 1:
	      ret = register_kret_new_func_1(func_name);
	      break;
	case 2:
	      ret = register_kret_new_func_2(func_name);
	      break;
	case 3:
	      ret = register_kret_new_func_3(func_name);
	      break;
	case 4:
	      ret = register_kret_new_func_4(func_name);
	      break;
	case 5:
	      ret = register_kret_new_func_5(func_name);
	      break;
	case 6:
	      ret = register_kret_new_func_6(func_name);
	      break;
	case 7:
	      ret = register_kret_new_func_7(func_name);
	      break;
	case 8:
	      ret = register_kret_new_func_8(func_name);
	      break;
	case 9:
	      ret = register_kret_new_func_9(func_name);
	      break;
	case 10:
	      ret = register_kret_new_func_10(func_name);
	      break;
	case 11:
	      ret = register_kret_new_func_11(func_name);
	      break;
	case 12:
	      ret = register_kret_new_func_12(func_name);
	      break;
	case 13:
	      ret = register_kret_new_func_13(func_name);
	      break;
	case 14:
	      ret = register_kret_new_func_14(func_name);
	      break;
	case 15:
	      ret = register_kret_new_func_15(func_name);
	      break;
	case 16:
	      ret = register_kret_new_func_16(func_name);
	      break;
	case 17:
	      ret = register_kret_new_func_17(func_name);
	      break;
	case 18:
	      ret = register_kret_new_func_18(func_name);
	      break;
	case 19:
	      ret = register_kret_new_func_19(func_name);
	      break;
	case 20:
	      ret = register_kret_new_func_20(func_name);
	      break;
	case 21:
	      ret = register_kret_new_func_21(func_name);
	      break;
	case 22:
	      ret = register_kret_new_func_22(func_name);
	      break;
	case 23:
	      ret = register_kret_new_func_23(func_name);
	      break;
	case 24:
	      ret = register_kret_new_func_24(func_name);
	      break;
	case 25:
	      ret = register_kret_new_func_25(func_name);
	      break;
	case 26:
	      ret = register_kret_new_func_26(func_name);
	      break;
	case 27:
	      ret = register_kret_new_func_27(func_name);
	      break;
	case 28:
	      ret = register_kret_new_func_28(func_name);
	      break;
	case 29:
	      ret = register_kret_new_func_29(func_name);
	      break;
	case 30:
	      ret = register_kret_new_func_30(func_name);
	      break;
	default:
	      break;
	}

	return ret;
}
void unregister_kret_new_func_batch(int index)
{
	switch(index) {
	case 1:
	      unregister_kret_new_func_1();
	      break;
	case 2:
	      unregister_kret_new_func_2();
	      break;
	case 3:
	      unregister_kret_new_func_3();
	      break;
	case 4:
	      unregister_kret_new_func_4();
	      break;
	case 5:
	      unregister_kret_new_func_5();
	      break;
	case 6:
	      unregister_kret_new_func_6();
	      break;
	case 7:
	      unregister_kret_new_func_7();
	      break;
	case 8:
	      unregister_kret_new_func_8();
	      break;
	case 9:
	      unregister_kret_new_func_9();
	      break;
	case 10:
	      unregister_kret_new_func_10();
	      break;
	case 11:
	      unregister_kret_new_func_11();
	      break;
	case 12:
	      unregister_kret_new_func_12();
	      break;
	case 13:
	      unregister_kret_new_func_13();
	      break;
	case 14:
	      unregister_kret_new_func_14();
	      break;
	case 15:
	      unregister_kret_new_func_15();
	      break;
	case 16:
	      unregister_kret_new_func_16();
	      break;
	case 17:
	      unregister_kret_new_func_17();
	      break;
	case 18:
	      unregister_kret_new_func_18();
	      break;
	case 19:
	      unregister_kret_new_func_19();
	      break;
	case 20:
	       unregister_kret_new_func_20();
	      break;
	case 21:
	      unregister_kret_new_func_21();
	      break;
	case 22:
	      unregister_kret_new_func_22();
	      break;
	case 23:
	      unregister_kret_new_func_23();
	      break;
	case 24:
	      unregister_kret_new_func_24();
	      break;
	case 25:
	      unregister_kret_new_func_25();
	      break;
	case 26:
	      unregister_kret_new_func_26();
	      break;
	case 27:
	      unregister_kret_new_func_27();
	      break;
	case 28:
	      unregister_kret_new_func_28();
	      break;
	case 29:
	      unregister_kret_new_func_29();
	      break;
	case 30:
	      unregister_kret_new_func_30();
	      break;
	default:
	      break;
	}
}

int register_kretprobe_ftrace = 0;
int register_ftrace_ftrace = 0;
int register_syscall_ftrace = 0;
int sysctl_kretprobe_enable_handler(struct ctl_table *table, int write,
		void __user *buffer, size_t *lenp, loff_t *ppos)
{
	int ret = 0;

	if (write && !capable(CAP_SYS_ADMIN))
		return -EPERM;

	if (sysctl_module_disable)
		return -EPERM;

	ret = proc_dostring(table, write, buffer, lenp, ppos);
	if (!write)
		return ret;

	if (register_kretprobe_ftrace == 1) {
		 pr_err("need unregister kretprobe function:maybe:%s", symbol_new);
		return -EPERM;
	}

	register_kretprobe_ftrace = 1;
	ret = register_kret_new_func(symbol_new);
	if (ret < 0) {
		register_kretprobe_ftrace = 0;
		strcpy(ftrace_hook_name, symbol_new);
		schedule_work(&ftrace_work_init);
	}
	return 0;
}

int sysctl_kretprobe_disable_handler(struct ctl_table *table, int write,
		void __user *buffer, size_t *lenp, loff_t *ppos)
{
	int ret;

	if (write && !capable(CAP_SYS_ADMIN))
		return -EPERM;

	if (sysctl_module_disable)
		return -EPERM;

	ret = proc_dostring(table, write, buffer, lenp, ppos);

	if (register_kretprobe_ftrace == 1)
		ungister__kret_new_func();

	schedule_work_on(0, &ftrace_work);

	register_kretprobe_ftrace = 0;

	return ret;
}

int syscall_hook_init(void)
{
	int ret;

	if (register_syscall_ftrace) {
	    pr_info("syscall hooked already");
	    return 0;
	}
	resident_kretprobe.kp.symbol_name = func_name;
	ret = register_kretprobe(&resident_kretprobe);
	if (ret < 0) {
		pr_err("register_kretprobe failed, returned %d\n", ret);
		return ret;
	}
	register_syscall_ftrace = 1;
	pr_info("Planted return probe at %s: %p\n",
			resident_kretprobe.kp.symbol_name, resident_kretprobe.kp.addr);
	return 0;
}

void  syscall_hook_exit(void)
{
	if (register_syscall_ftrace == 0)
	    return;

	register_syscall_ftrace = 0;

	unregister_kretprobe(&resident_kretprobe);
	pr_info("kretprobe at %p unregistered\n", resident_kretprobe.kp.addr);

	/* nmissed > 0 suggests that maxactive was set too low. */
	pr_info("Missed probing %d instances of %s\n",
		resident_kretprobe.nmissed, resident_kretprobe.kp.symbol_name);
}

int sysctl_system_hook_handler(struct ctl_table *table, int write,
		void __user *buffer, size_t *lenp, loff_t *ppos)
{
	int ret;
	if (write && !capable(CAP_SYS_ADMIN))
		return -EPERM;

	if (sysctl_module_disable)
		return -EPERM;

	ret = proc_dointvec_minmax(table, write, buffer, lenp, ppos);

	if (!write)
        return ret;

	syscall_hook_init();
	system_base_function_hook();

	return ret;
}
int sysctl_system_unhook_handler(struct ctl_table *table, int write,
		void __user *buffer, size_t *lenp, loff_t *ppos)
{
	int ret;

	if (write && !capable(CAP_SYS_ADMIN))
		return -EPERM;

	if (sysctl_module_disable)
		return -EPERM;

	ret = proc_dointvec_minmax(table, write, buffer, lenp, ppos);

	if (!write)
        return ret;

	syscall_hook_exit();
	system_base_function_unhook();
	return ret;
}

int kretprobe_init(void)
{
	return 0;
}

void  kretprobe_exit(void)
{

	if (register_kretprobe_ftrace == 1)
		ungister__kret_new_func();

	while (register_ftrace_ftrace == 3)
		schedule_timeout(1000);

	ftrace_unhook_work_fn(&ftrace_work);

	if (register_syscall_ftrace == 1)
		syscall_hook_exit();

	register_kretprobe_ftrace = 0;
	register_syscall_ftrace = 0;
}

